import torch.nn
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
import resnet50

transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize(mean=[0.5,0.5,0.5],std=[0.5,0.5,0.5]),
                                transforms.Resize((224, 224))
                               ])

training_data = datasets.CIFAR10(root="data",train=True,download=True,transform=transform)

testing_data = datasets.CIFAR10(root="data",train=False,download=True,transform=transform)

batch_size = 64
train_data = DataLoader(dataset=training_data,batch_size=batch_size,shuffle=True,drop_last=True)
test_data = DataLoader(dataset=testing_data,batch_size=batch_size,shuffle=True,drop_last=True)


epochs = 10
learning_rate = 0.01

model = resnet50.ResNet50(resnet50.Bottleneck)
loss = torch.nn.CrossEntropyLoss()
optim = torch.optim.SGD(model.parameters(), lr=learning_rate)


for epoch in range(epochs):
    model.train()
    for data in train_data:
        imgs, targets = data
        outputs = model(imgs)
        result_loss = loss(outputs, targets)
        optim.zero_grad()
        result_loss.backward()
        optim.step()